import numpy as np
import matplotlib.pyplot as plt
from tqdm import trange
import copy


plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
plt.rc('font', size=10)
plt.rc('font', family='SimHei')


# 通用常数
interval = 1  # 投弹时间间隔
v_down = 3  # 云团下降速度
r_effective = 10  # 有效遮蔽半径
t_effective = 20  # 有效遮蔽时间
v_missile = 300  # 导弹速度
r_target = 7  # 目标半径
H_target = 10  # 目标高度
target = np.array([0, 200, 0])  # 目标底部圆心坐标
M1 = np.array([20000, 0, 2000])
M2 = np.array([19000, 600, 2100])
M3 = np.array([18000, -600, 1900])
FY1 = np.array([17800, 0, 1800])
FY2 = np.array([12000, 1400, 1400])
FY3 = np.array([6000, -3000, 700])
FY4 = np.array([11000, 2000, 1800])
FY5 = np.array([13000, -2000, 1300])
v_range = np.array([70, 140])  # 无人机速度范围
g = 9.8  # 重力加速度
target_point = np.array([
    [0, 200, 0]
])
t_total = np.linalg.norm(M1) / 300  # 从0开始到导弹击中目标的总用时
precision = 0.01  # 时间精度


# 传入M1的初始坐标（也即可确定其飞行方向）和经过的时间，返回其在给定时间下的坐标
def get_M1_coordinate(coordinate_0, t, v=None):
    if not v:
        v = -coordinate_0 / np.linalg.norm(coordinate_0) * v_missile
    else:
        v = v / np.linalg.norm(v) * v_missile
    return coordinate_0 + v * t


# 传入FY1的坐标和经过的时间，返回它投射的第i枚烟雾弹中心的坐标
def get_FY1_coordinate(coordinate_0, t, i, v_1, t_1, t_2, v=None):
    t_drop = 0
    for j in range(i + 1):
        t_drop += t_1[i]
    t_explosion = t_2[i]
    if t < t_drop + t_explosion or t > t_drop + t_explosion + t_effective:
        return [None, None, None]
    else:
        if v is None:
            v = -coordinate_0
        else:
            v = v - coordinate_0
        v[2] = 0
        v = v / np.linalg.norm(v) * v_1
        delta_t = t - t_drop - t_explosion
        coordinate = coordinate_0 + v * (t_drop + t_explosion)
        coordinate[2] -= 1 / 2 * g * (t_explosion ** 2)
        coordinate[2] -= delta_t * v_down
        return coordinate


def collision_detection(coordinate_missile, coordinate_0, t, v_1, t_1, t_2, v=None):
    global target_point, r_effective
    count_collision = 0
    for target in target_point:
        for i in range(3):
            coordinate_bomb = get_FY1_coordinate(coordinate_0, t, i, v_1, t_1, t_2, v)
            if coordinate_bomb[0] == None:
                break
            l_1 = coordinate_bomb - target
            l_2 = coordinate_missile - target
            molecule = np.cross(l_1, l_2)
            denominator = np.linalg.norm(l_2)
            d = np.linalg.norm(molecule) / denominator
            if d < r_effective:
                missile_to_target = -coordinate_missile  # 由导弹指向目标
                missile_to_bomb = coordinate_bomb - coordinate_missile  # 由导弹指向烟幕弹
                cos = np.dot(missile_to_target, missile_to_bomb)
                if cos > 0 or np.linalg.norm(missile_to_bomb) < r_effective:  # 余弦值大于0，是锐角，或者导弹和烟幕弹之间的距离小于有效半径
                    count_collision += 1
                    break
    return count_collision


# 飞行方向、飞行速度、烟幕干扰弹投放点、烟幕干扰弹起爆点
def fun(v_direction_x, v_diection_y, v_1, t_1, t_2):
    t_list = np.arange(t_1[0] + t_2[0], t_total, precision)
    count = 0
    for t in t_list:
        if collision_detection(get_M1_coordinate(M1, t), FY1, t, v_1, t_1, t_2, np.array([v_direction_x, v_diection_y, 0])) == len(target_point):
            count += 1
    return count


# 在SA函数外部初始化记录容器
history = {
    'best_P': [],
    'current_P': [],
    'k': [[], [], [], [], [], [], [], [], []],
}


def plot_SA_history():
    plt.figure(figsize=(15, 8))

    # 子图1：损失函数变化
    plt.subplot(2, 3, 1)
    plt.plot(history['best_P'], 'r-', label='Best P')
    plt.plot(history['current_P'], 'b--', alpha=0.5, label='Current P')
    plt.xlabel('Iteration')
    plt.ylabel('P')
    plt.title('P变化曲线')
    plt.legend()
    plt.grid(True, which="both", ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 2)
    plt.plot(history['k'][0], label="目标点x坐标分量")
    plt.plot(history['k'][1], label="目标点y坐标分量")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('目标点参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 3)
    plt.plot(history['k'][2], label="无人机速度大小")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('飞行速度变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 4)
    plt.plot(history['k'][3], label="投弹时间")
    plt.plot(history['k'][4], label="爆炸时间")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('烟幕弹1时间参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 5)
    plt.plot(history['k'][5], label="投弹时间")
    plt.plot(history['k'][6], label="爆炸时间")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('烟幕弹2时间参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 6)
    plt.plot(history['k'][7], label="投弹时间")
    plt.plot(history['k'][8], label="爆炸时间")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('烟幕弹3时间参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    plt.tight_layout()
    plt.show()


def SA(iter, t0, tf, alpha):
    t = t0
    v_direction_x_c = 0
    v_diection_y_c = 200
    v_1_c = 120
    t_1_c = [1.5, 1.5, 1.5]
    t_2_c = [3.6, 3.6, 3.6]
    Pc = fun(v_direction_x_c, v_diection_y_c, v_1_c, t_1_c, t_2_c)

    v_direction_x_b = v_direction_x_c
    v_diection_y_b = v_diection_y_c
    v_1_b = v_1_c
    t_1_b = [i for i in t_1_c]
    t_2_b = [i for i in t_2_c]
    Pb = Pc
    for i in trange(iter):
        v_direction_x_n = v_direction_x_c + np.random.normal(0, 50)
        v_diection_y_n = v_diection_y_c + np.random.normal(0, 50)
        v_1_n = v_1_c + np.random.normal(0, 2)
        t_1_n = [
            t_1_c[0] + np.random.normal(0, 0.1),
            t_1_c[1] + np.random.normal(0, 0.1),
            t_1_c[2] + np.random.normal(0, 0.1),
        ]
        t_2_n = [
            t_2_c[0] + np.random.normal(0, 0.5),
            t_2_c[1] + np.random.normal(0, 0.1),
            t_2_c[2] + np.random.normal(0, 0.1),
        ]

        v_direction_x_n = np.clip(v_direction_x_n, -300, 300)
        v_diection_y_n = np.clip(v_diection_y_n, 0, 450)
        v_1_n = np.clip(v_1_n, v_range[0], v_range[1])
        t_1_n[0] = np.clip(t_1_n[0], 0, t_total)
        t_1_n[1] = np.clip(t_1_n[1], 1, 3)
        t_1_n[2] = np.clip(t_1_n[2], 1, 3)
        t_2_n[0] = np.clip(t_2_n[0], 0, 5)
        t_2_n[1] = np.clip(t_2_n[1], 0, 5)
        t_2_n[2] = np.clip(t_2_n[2], 0, 5)
        Pn = fun(v_direction_x_n, v_diection_y_n, v_1_n, t_1_n, t_2_n)
        if Pn > Pc or np.random.rand() < np.exp((Pn-Pc)/t):
            v_direction_x_c = v_direction_x_n
            v_diection_y_c = v_diection_y_n
            v_1_c = v_1_n
            t_1_c = [i for i in t_1_n]
            t_2_c = [i for i in t_2_n]
            Pc = Pn
            if Pc > Pb:
                v_direction_x_b = v_direction_x_c
                v_diection_y_b = v_diection_y_c
                v_1_b = v_1_c
                t_1_b = [i for i in t_1_c]
                t_2_b = [i for i in t_2_c]
                Pb = Pc
        t = t * alpha
        if t < tf:
            break
        print(f"第{i}轮,Pb:{Pb},Pc:{Pc},kb:{(v_direction_x_b, v_diection_y_b, v_1_b, t_1_b, t_2_b)},kc:{(v_direction_x_c, v_diection_y_c, v_1_c, t_1_c, t_2_c)}")
        history['best_P'].append(Pb)
        history['current_P'].append(Pc)
        history['k'][0].append(v_direction_x_b)
        history['k'][1].append(v_diection_y_b)
        history['k'][2].append(v_1_b)
        history['k'][3].append(t_1_b[0])
        history['k'][4].append(t_2_b[0])
        history['k'][5].append(t_1_b[0] + t_1_b[1])
        history['k'][6].append(t_2_b[1])
        history['k'][7].append(t_1_b[0] + t_1_b[1] + t_1_b[2])
        history['k'][8].append(t_2_b[2])
    print(f"无人机FY1从初始位置向着点{v_direction_x_b, v_diection_y_b, 0}飞, 速度为{v_1_b}m/s,开始投第一颗烟幕弹时刻为{t_1_b[0]}s, 爆炸间隔为{t_2_b[0]}s; 开始投第二颗烟幕弹时刻为{t_1_b[0] + t_1_b[1]}s, 爆炸间隔为{t_2_b[1]}s; 开始投第三颗烟幕弹时刻为{t_1_b[0] + t_1_b[1] + t_1_b[2]}s, 爆炸间隔为{t_2_b[2]}s,最终有效遮挡时间为{Pb * precision}s")
    return (v_direction_x_b, v_diection_y_b, v_1_b, t_1_b, t_2_b), Pb


re = SA(150, 800, 0.001, 0.95)
print(re)
plot_SA_history()
